{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QKWvMXWMBEJA"
      },
      "source": [
        "# Behavioural Planning for Autonomous Highway Driving\n",
        "\n",
        "We plan a trajectory using the _Optimistic Planning for Deterministic systems_ ([OPD](https://hal.inria.fr/hal-00830182)) algorithm.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "s-ghXis6A_md"
      },
      "outputs": [],
      "source": [
        "#@title Imports for env, agent, and visualisation.\n",
        "# Environment\n",
        "!pip install highway-env\n",
        "import gymnasium as gym\n",
        "import highway_env\n",
        "\n",
        "# Agent\n",
        "!pip install git+https://github.com/eleurent/rl-agents#egg=rl-agents\n",
        "from rl_agents.agents.common.factory import agent_factory\n",
        "\n",
        "# Visualisation\n",
        "import sys\n",
        "from tqdm.notebook import trange\n",
        "!pip install moviepy -U\n",
        "!pip install imageio_ffmpeg\n",
        "!pip install pyvirtualdisplay\n",
        "!apt-get install -y xvfb ffmpeg\n",
        "!git clone https://github.com/Farama-Foundation/HighwayEnv.git\n",
        "sys.path.insert(0, './highway-env/scripts/')\n",
        "from utils import record_videos, show_videos\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bgNDDWwqCj8l"
      },
      "outputs": [],
      "source": [
        "#@title Run an episode\n",
        "\n",
        "# Make environment\n",
        "env = gym.make(\"highway-fast-v0\", render_mode=\"rgb_array\")\n",
        "env = record_videos(env)\n",
        "(obs, info), done = env.reset(), False\n",
        "\n",
        "# Make agent\n",
        "agent_config = {\n",
        "    \"__class__\": \"<class 'rl_agents.agents.tree_search.deterministic.DeterministicPlannerAgent'>\",\n",
        "    \"env_preprocessors\": [{\"method\":\"simplify\"}],\n",
        "    \"budget\": 50,\n",
        "    \"gamma\": 0.7,\n",
        "}\n",
        "agent = agent_factory(env, agent_config)\n",
        "\n",
        "# Run episode\n",
        "for step in trange(env.unwrapped.config[\"duration\"], desc=\"Running...\"):\n",
        "    action = agent.act(obs)\n",
        "    obs, reward, done, truncated, info = env.step(action)\n",
        "    \n",
        "env.close()\n",
        "show_videos()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "highway-planning.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "pycharm": {
      "stem_cell": {
        "cell_type": "raw",
        "metadata": {
          "collapsed": false
        },
        "source": []
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
